# -*- coding: utf-8 -*-
import socket
import threading
import time

class TCPSocketCommunicator:
    """
    TCP Socket 通信类，支持作为客户端连接或作为服务器监听，并进行字符串收发。
    """
    def __init__(self, local_ip: str, local_port: int, buffer_size: int = 4096):
        """
        初始化 TCP Socket 通信器。

        Args:
            local_ip (str): 本地绑定的 IP 地址。使用 '0.0.0.0' 监听所有可用接口。
            local_port (int): 本地绑定的端口号。
            buffer_size (int): 接收缓冲区大小，默认为 4096 字节。
        """
        self.local_ip = local_ip
        self.local_port = local_port
        self.buffer_size = buffer_size
        self.socket = None  # 用于监听或连接的主 socket
        self.connection = None  # 建立连接后的客户端 socket (服务器端) 或 连接 socket (客户端)
        self.addr = None    # 连接对方的地址
        self._is_connected = False
        self._is_listening = False
        self._lock = threading.Lock() # 用于线程安全的发送

        print(f"TCPSocketCommunicator initialized for {local_ip}:{local_port}")

    def _create_socket(self):
        """创建 TCP socket。"""
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            # 允许端口重用，防止 TIME_WAIT 状态导致绑定失败
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            return sock
        except Exception as e:
            print(f"Error creating socket: {e}")
            return None

    def start_listening(self, max_connections: int = 1):
        """
        作为服务器启动监听模式，等待客户端连接。
        此方法会阻塞，直到有客户端连接。

        Args:
            max_connections (int): 最大等待连接数。
        
        Returns:
            bool: 如果成功建立连接则返回 True，否则返回 False。
        """
        if self._is_connected or self._is_listening:
            print("Already connected or listening.")
            return False

        self.socket = self._create_socket()
        if not self.socket:
            return False

        try:
            self.socket.bind((self.local_ip, self.local_port))
            self.socket.listen(max_connections)
            self._is_listening = True
            print(f"Listening on {self.local_ip}:{self.local_port}...")

            self.connection, self.addr = self.socket.accept()
            self._is_connected = True
            print(f"Accepted connection from {self.addr}")
            return True
        except Exception as e:
            print(f"Error starting listening: {e}")
            self._cleanup()
            return False

    def connect_to_server(self, remote_ip: str, remote_port: int, timeout: float = 5.0):
        """
        作为客户端连接到远程服务器。

        Args:
            remote_ip (str): 远程服务器的 IP 地址。
            remote_port (int): 远程服务器的端口号。
            timeout (float): 连接超时时间（秒）。

        Returns:
            bool: 如果成功连接则返回 True，否则返回 False。
        """
        if self._is_connected or self._is_listening:
            print("Already connected or listening.")
            return False

        self.socket = self._create_socket()
        if not self.socket:
            return False

        self.socket.settimeout(timeout) # 设置连接超时
        try:
            print(f"Connecting to {remote_ip}:{remote_port}...")
            self.socket.connect((remote_ip, remote_port))
            self.connection = self.socket # 客户端模式下，connection 就是主 socket
            self.addr = (remote_ip, remote_port)
            self._is_connected = True
            print(f"Successfully connected to {self.addr}")
            return True
        except socket.timeout:
            print(f"Connection to {remote_ip}:{remote_port} timed out.")
            self._cleanup()
            return False
        except Exception as e:
            print(f"Error connecting to server: {e}")
            self._cleanup()
            return False
        finally:
            self.socket.settimeout(None) # 恢复默认阻塞模式

    def send_message(self, message: str):
        """
        发送字符串消息。

        Args:
            message (str): 要发送的字符串。

        Returns:
            bool: 如果发送成功则返回 True，否则返回 False。
        """
        if not self._is_connected or self.connection is None:
            print("Not connected. Cannot send message.")
            return False

        try:
            with self._lock: # 确保发送操作的线程安全
                self.connection.sendall(message.encode('utf-8'))
            return True
        except Exception as e:
            print(f"Error sending message: {e}")
            self._cleanup() # 发送失败通常意味着连接断开，需要清理
            return False

    def receive_message(self):
        """
        接收字符串消息。

        Returns:
            str: 接收到的字符串消息，如果连接断开或接收失败则返回 None。
        """
        if not self._is_connected or self.connection is None:
            print("Not connected. Cannot receive message.")
            return None

        try:
            data = self.connection.recv(self.buffer_size)
            if not data:
                print("Connection closed by remote host.")
                self._cleanup()
                return None
            return data.decode('utf-8')
        except Exception as e:
            print(f"Error receiving message: {e}")
            self._cleanup() # 接收失败通常意味着连接断开，需要清理
            return None

    def is_connected(self):
        """
        检查当前是否已建立连接。
        """
        return self._is_connected

    def _cleanup(self):
        """
        内部清理方法，关闭所有 socket 连接。
        """
        if self.connection:
            try:
                self.connection.shutdown(socket.SHUT_RDWR) # 关闭读写
                self.connection.close()
            except OSError as e:
                # 忽略 socket 已经关闭的错误
                if e.errno != 107: # errno 107 is ENOTCONN (Transport endpoint is not connected)
                    print(f"Error during connection shutdown/close: {e}")
            self.connection = None
        
        if self.socket:
            try:
                self.socket.close()
            except Exception as e:
                print(f"Error during main socket close: {e}")
            self.socket = None

        self._is_connected = False
        self._is_listening = False
        print("Socket resources cleaned up.")

    def close(self):
        """
        关闭通信器，释放所有资源。
        """
        self._cleanup()

# --- 示例用法 ---
if __name__ == "__main__":
    # 示例 1: 作为服务器运行
    def run_server():
        server_ip = '0.0.0.0' # 监听所有可用接口
        server_port = 12345
        server_comm = TCPSocketCommunicator(server_ip, server_port)
        
        try:
            if server_comm.start_listening():
                print("Server ready to receive messages...")
                while server_comm.is_connected():
                    received_msg = server_comm.receive_message()
                    if received_msg:
                        print(f"Server received: {received_msg}")
                        if received_msg.lower() == "exit":
                            print("Server exiting due to 'exit' command.")
                            break
                        server_comm.send_message(f"Echo: {received_msg}")
                    else:
                        # 连接已关闭或接收失败
                        break
        except KeyboardInterrupt:
            print("\nServer interrupted by user.")
        finally:
            server_comm.close()
            print("Server closed.")

    # 示例 2: 作为客户端运行
    def run_client():
        client_ip = '0.0.0.0' # 客户端绑定到任意本地接口
        client_port = 0 # 让系统自动分配一个可用端口
        remote_ip = '127.0.0.1' # 假设服务器在本地运行
        remote_port = 12345
        client_comm = TCPSocketCommunicator(client_ip, client_port)

        try:
            if client_comm.connect_to_server(remote_ip, remote_port):
                messages_to_send = ["Hello Server!", "How are you?", "This is a test.", "exit"]
                for msg in messages_to_send:
                    print(f"Client sending: {msg}")
                    if client_comm.send_message(msg):
                        received_msg = client_comm.receive_message()
                        if received_msg:
                            print(f"Client received: {received_msg}")
                        else:
                            print("Client failed to receive response or server closed connection.")
                            break
                    else:
                        print("Client failed to send message.")
                        break
                    time.sleep(1) # 等待一下
            else:
                print("Client failed to connect.")
        except KeyboardInterrupt:
            print("\nClient interrupted by user.")
        finally:
            client_comm.close()
            print("Client closed.")

    # 你可以同时运行这两个函数在不同的终端中进行测试
    # 例如：在一个终端运行 run_server()，在另一个终端运行 run_client()

    print("Choose mode: (s)erver or (c)lient")
    mode = input().lower()

    if mode == 's':
        run_server()
    elif mode == 'c':
        run_client()
    else:
        print("Invalid mode selected. Please choose 's' or 'c'.")

